#
# Copyright (C) 2023, Inria
# GRAPHDECO research group, https://team.inria.fr/graphdeco
# All rights reserved.
#
# This software is free for non-commercial, research and evaluation use 
# under the terms of the LICENSE.md file.
#
# For inquiries contact  george.drettakis@inria.fr
#

from typing import List

import torch
from torch import Tensor
import numpy as np
from numpy import ndarray

from utils.graphics_utils import getWorld2View2, getProjectionMatrix


class Camera:

    def __init__(self, uid:int, colmap_id:int, R:ndarray, T:ndarray, FoVx:float, FoVy:float, 
                 gt_image:Tensor, mask:Tensor=None, freq_images:List[Tensor]=[], image_name:str=None, trans:ndarray=np.zeros([3]), scale:float=1.0,
                 data_device:str='cuda'):
        self.uid = uid
        self.colmap_id = colmap_id

        self.R = R
        self.T = T
        self.FoVx = FoVx
        self.FoVy = FoVy
        self.zfar = 100.0
        self.znear = 0.01
        self.trans = trans
        self.scale = scale
        self.freq_images = freq_images

        self.world_view_transform = torch.tensor(getWorld2View2(R, T, trans, scale)).transpose(0, 1).cuda()
        self.projection_matrix = getProjectionMatrix(znear=self.znear, zfar=self.zfar, fovX=self.FoVx, fovY=self.FoVy).transpose(0,1).cuda()
        self.full_proj_transform = (self.world_view_transform.unsqueeze(0).bmm(self.projection_matrix.unsqueeze(0))).squeeze(0)
        self.camera_center = self.world_view_transform.inverse()[3, :3]

        img = gt_image.clamp_(0.0, 1.0)
        if mask is not None:
            img *= mask.to(device=img)
        self.gt_image = img
        self.image_name = image_name
        self.image_height = self.gt_image.shape[1]
        self.image_width = self.gt_image.shape[2]


class MiniCam:
    def __init__(self, width, height, fovy, fovx, znear, zfar, world_view_transform, full_proj_transform):
        self.image_width = width
        self.image_height = height    
        self.FoVy = fovy
        self.FoVx = fovx
        self.znear = znear
        self.zfar = zfar
        self.world_view_transform = world_view_transform
        self.full_proj_transform = full_proj_transform
        view_inv = torch.inverse(self.world_view_transform)
        self.camera_center = view_inv[3][:3]
